#include "kernel_operator.h"
using namespace AscendC;
constexpr uint32_t BUFFER_NUM = 1;
#include "sum.cpp"
#include "mean.cpp"


extern "C" __global__ __aicore__ void nll_loss(GM_ADDR x, GM_ADDR target, GM_ADDR weight, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling){
    GET_TILING_DATA(tiling_data, tiling);
    TPipe pipe;
    uint32_t bn;
    if(GetBlockIdx()<tiling_data.fn){
        bn = tiling_data.fb;
    }else{
        bn = tiling_data.tb;
    }
    if(TILING_KEY_IS(0)){
        KernelNLLLossSum op;
        op.Init(&pipe, x, target, weight, y, workspace,
            tiling_data.N, tiling_data.C, bn, tiling_data.fb
            , tiling_data.fn
        ); 
        op.Process();
    }
    if(TILING_KEY_IS(1)){
        KernelNLLLossMean op;
        op.Init(&pipe, x, target, weight, y, workspace,
            tiling_data.N, tiling_data.C, bn, tiling_data.fb
            , tiling_data.fn
        ); 
        op.Process();
    }
}